#pragma once
#ifndef AVGPOOL_H_
#define AVGPOOL_H_

#include "stream_tools.h"

template<
	unsigned N_IO,		// # of IO
	unsigned N_CH,		// # of CHannel
	unsigned HIGHT,		// # of HIGHT
	unsigned WIDTH,		// # of WIDTH
	unsigned N_BTACH	// # of BTACH
>
void avgpool(
	data_stream<N_IO * BIT_ACTV>& in, 
	data_stream<N_IO * BIT_CONV>& out
)
{
	static_assert(N_CH >= N_IO, "avgpool");
	static_assert(N_CH % N_IO == 0, "avgpool");
	constexpr unsigned FOLD = N_CH / N_IO;
	constexpr unsigned ITERS = HIGHT*WIDTH;

	assert(in.size() == HIGHT*WIDTH*N_CH/N_IO*N_BTACH);
	assert(out.empty());

	#pragma HLS DATAFLOW

	ap_uint<BIT_CONV> acc[FOLD][N_IO];
	#pragma HLS ARRAY_PARTITION variable=acc complete dim=2

	for (unsigned f = 0; f < FOLD; ++f)
	{
	#pragma HLS PIPELINE II=1
		for (unsigned i = 0; i < N_IO; ++i)
		{
		#pragma HLS UNROLL
			acc[f][i] = 0;
		}
	}

	for (unsigned n = 0; n < N_BTACH; ++n)
	{
		for (unsigned it = 0; it < ITERS; ++it)
		{
			for (unsigned f = 0; f < FOLD; ++f)
			{
			#pragma HLS PIPELINE II=1
				// load
				ap_uint<N_IO * BIT_ACTV> in_buf = in.read();
				// calc
				for (unsigned i = 0; i < N_IO; ++i)
				{
					#pragma HLS UNROLL
					ap_uint<BIT_ACTV> x = in_buf(SLICE(BIT_ACTV, i));
					acc[f][i] += x;
				}
			}
		}

		for (unsigned f = 0; f < FOLD; ++f)
		{	
		#pragma HLS PIPELINE II=1
			ap_uint<N_IO * BIT_CONV> out_buf;
			for (unsigned i = 0; i < N_IO; ++i)
			{
			#pragma HLS UNROLL
				out_buf(SLICE(BIT_CONV, i)) = acc[f][i];
				acc[f][i] = 0;
			}
			out.write(out_buf);
		}
	}
	assert(in.empty());
	assert(out.size() == N_CH*N_BTACH);
	return;
}

#endif
